# ====================================
# Company: Zonge International, INC.
# Developer: Wanjie Feng
# Date: 10/18/2023
# Time: 9:54 AM
# Filename: ERT_Process.py
# IDE: PyCharm
# =====================================
import numpy as np
import matplotlib.pyplot as plt
import time
from matplotlib import cm
from scipy.interpolate import griddata
import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.patches import Rectangle
import io
'''
    本程序模拟了手部UI读取数据、显示数据的过程。这里假设主机传到手部的文件格式为*.cvs格式，
    如后期数据格式有变，数据读取部分也要做相应改变。
    1. 获取硬件信息
    2. 项目参数
    3. 读取 *.csv 数据及部分系统状态信息
    4. 形成头文件
    5. 处理数据
    6. 显示数据
    7. 保存数据
    8. 重复3-7
    9. 采集结束
'''

from Read_System import ReadSys
from Read_Survey import ReadSurvey
from Data_collection import ReadData
from Compute_RhoIP import ComputeRhoIP
from ERT_Plot import ERTPlot

# 排列图示
def ERTuiArray(num_electrodes0,Max_n_spacing0,Aspacing0,array_type0):

        # 开始采集之前，先计算排列，这里的参数不需要保存到数据文件里。
        #num_electrodes 电极总数（来自主机）
        #Max_n_spacing 最大隔离系数，若不设置，会根据电极总数计算，为 (num_electrodes - 3) （手部输入）
        #Aspacing 点距 （这个手部输入）
        #array_type 排列方式
        print("------------------ERTuiArray--------------------")
        num_electrodes = num_electrodes0
        Max_n_spacing = Max_n_spacing0
        Aspacing = Aspacing0
        array_type = array_type0
        print(num_electrodes)
        print(Max_n_spacing)
        print(Aspacing)
        print(array_type)
        # 最大隔离系数太大，系统自动调整
        if Max_n_spacing > num_electrodes - 3:
            print('最大隔离系数设置错误，隔离系数范围为' + str(1) + '到' + str(num_electrodes - 3))
            Max_n = num_electrodes - 3

        NS = np.sum(range(1, Max_n_spacing + 1)) + Max_n_spacing * (
                    (num_electrodes - 3) - Max_n_spacing)  # 总排列数

        ABMN = np.zeros((NS, 4))  # ABMN电极排列
        c = 0
        # 计算排列电极位置 （这里假设只有一个ADC接收电极，当有多个ADC时，此部分再修改)
        if array_type == 0: #dipole-dipole
            for i in range(num_electrodes - 3):
                for j in range(i + 2, num_electrodes - 1):
                    ABMN[c, 0] = i + 1  # A电极
                    ABMN[c, 1] = i + 2  # B电极
                    ABMN[c, 2] = j + 1  # M电极
                    ABMN[c, 3] = j + 2  # N电极
                    N = ABMN[c, 2] - ABMN[c, 1]
                    if N > Max_n_spacing:
                        continue
                    else:
                        c = c + 1
        elif array_type == 1:#pole-dipole
            for i in range(num_electrodes - 3):
                for j in range(i + 2, num_electrodes - 1):
                    ABMN[c, 0] = 1  # A电极, 整个测量中不动
                    ABMN[c, 1] = i + 2  # B电极
                    ABMN[c, 2] = j + 1  # M电极
                    ABMN[c, 3] = j + 2  # N电极
                    N = ABMN[c, 2] - ABMN[c, 1]
                    if N > Max_n_spacing:
                        continue
                    else:
                        c = c + 1
        ###3. 读取CVS数据及系统参数
        '''
            手部首先需要向主机发送数据采集的命令，规定采集的参数，然后主机才能返回数据。这里假设命令发送成果，数据成功上传。
            数据按每次叠加上传，例如我们设置命令为叠加5次，每次波形长度为8，那么手部会收到8个完整波形后再做相应的处理。
            这8个波形的数据在数据文件中是一个数据块，有相应的头文件与之对应。
        '''
        IPValue = {}
        # for AI in range(NS):
        #     print("------------------ai-------------------------")
        #     print(AI)
        print("------------------ai-------------------------")
        print(print(ABMN[0,:]))
        IPValue[0] = ComputeRhoIPArray(ABMN[0,:],Aspacing)

            ###4. 画出图形
        return ERTuiArrayShow(IPValue, ABMN, Aspacing)
        # print("===================================开始进入ERT_PLOT============================================")
def ComputeRhoIPArray(ABMN,Aspacing):
    IPValues = {}
    TXA = (ABMN[0] - 1) * Aspacing #Aspacing 点距 也是app 填写的
    TXB = (ABMN[1] - 1) * Aspacing
    RXM = (ABMN[2] - 1) * Aspacing
    RXN = (ABMN[3] - 1) * Aspacing

    IPValues['TXA'] = TXA
    IPValues['TXB'] = TXB
    IPValues['RXM'] = RXM
    IPValues['RXN'] = RXN
    IPValues['X'] = (RXM-TXB)/2 + TXB                           #画图横坐标位置
    IPValues['Z'] = (RXM-TXB)/2 + Aspacing/2      #画图垂向坐标位置
    IPValues['n'] = ABMN[2] - ABMN[1]                           #隔离系数
    return IPValues
def ERTuiArrayShow(IPValues, ABMN, Aspacing):

    '''
    本程序画出高密度电阻率法的等值线图和排列图
    :param IPValues:
    :return: NONE
'''
    ####这里画出两个图，一个是排列图，利用红色圆圈表示发射，用黑色圆表示接收
    #最后一个IP的值
    #获取相应的画图参数
    #Aspacing 点距

    NIP = len(IPValues)
    lastIP = IPValues[NIP-1]
    print("------------- lastIP ------------"+str(lastIP))
    #print("------------- IPValues ------------"+str(IPValues))

    #print("------------- ABMN ------------"+str(ABMN))

    #print("------------- Aspacing ------------"+str(Aspacing))

    #print("------------- NIP ------------"+str(NIP))

    #print("------------- lastIP ------------"+str(lastIP))
    TXA = lastIP['TXA']
    TXB = lastIP['TXB']
    RXM = lastIP['RXM']
    RXN = lastIP['RXN']
    X = lastIP['X']
    Z = lastIP['Z']
    #rho = lastIP['Rho']
    # #####计算图的尺寸
    Min_X = np.min(ABMN[:,0]-1) * Aspacing
    Max_X = np.max(ABMN[:,3]-1) * Aspacing
    Min_Z = np.min(ABMN[:,2]-ABMN[:,1])*Aspacing/2 + Aspacing/2
    Max_Z = np.max(ABMN[:,2]-ABMN[:,1])*Aspacing/2 + Aspacing/2
    print("------------- Min_X ------------"+str(Min_X))
    print("------------- Max_X ------------"+str(Max_X))



    fig1, ax1 = plt.subplots()
    fig1.patch.set_alpha(0.)  # 或者其他值
    plt.figure(figsize=(24,8))
    ax1.clear()                                     #Python清除上次画图的内容
    plt.xlim([Min_X,Max_X])                         #设置横向尺度
    plt.ylim([0,Max_Z])
    plt.figure(figsize=(22,10))
    #设置纵向尺度
    Radius = Aspacing/10                            #圆表示画图点，圆的半径
    #画出所有排列的地下记录点
    NS = np.size(ABMN,0)
    for CI in range(NS):
        BE = (ABMN[CI, 1]-1) * Aspacing
        ME = (ABMN[CI, 2]-1) * Aspacing
        XE = (ME-BE)/2 + BE
        ZE = (ME-BE)/2 + Aspacing/2
        E_C = plt.Circle((XE, ZE), radius=Radius, fill=False, color='black', linewidth=1.0, \
                         clip_on=False)  # 画图点位置
        plt.gca().add_artist(E_C)
    #画出本次测量的地下记录的及发射接收位置
    TXA_C = plt.Circle((TXA, 0), radius=Radius, fill=False, color='red', linewidth=2.0, \
                       clip_on=False)               #发射A位置
    plt.gca().add_artist(TXA_C)
    TXB_C = plt.Circle((TXB, 0), radius=Radius, fill=False, color='red', linewidth=2.0, \
                       clip_on=False)               #发射B位置
    plt.gca().add_artist(TXB_C)
    RXM_C = plt.Circle((RXM, 0), radius=Radius, fill=False, color='black', linewidth=2.0, \
                       clip_on=False)               #接收M位置
    plt.gca().add_artist(RXM_C)
    RXN_C = plt.Circle((RXN, 0), radius=Radius, fill=False, color='black', linewidth=2.0, \
                       clip_on=False)               #接收N位置
    plt.gca().add_artist(RXN_C)
    #电阻率画图点用实心圆表示
    Rho_C = plt.Circle((X, Z), radius=Radius, fill=True, color='blue', linewidth=2.0, \
                       clip_on=False)               #拟断面图画图点
    plt.gca().add_artist(Rho_C)
    plt.gca().invert_yaxis()
    plt.gca().set_xlabel('Array Location(m)', size=14)
    plt.gca().set_ylabel('n-spacing', size=14)
    plt.gca().xaxis.set_label_position('top')
    plt.gca().xaxis.tick_top()
    plt.gca().set_aspect('equal')
    plt.yticks(np.insert(np.arange(Min_Z, Max_Z + Aspacing/2, Aspacing/2),0,0))
    plt.xticks(np.arange(Min_X, Max_X, Aspacing ))
    #用隔离系数代替纵向坐标的标签
    labels = [item.get_text() for item in ax1.get_yticklabels()]
    NLAB = len(labels)
    labs = np.insert(np.arange(int(np.min(ABMN[:, 2] - ABMN[:, 1])), \
                               int(np.max(ABMN[:, 2] - ABMN[:, 1]) + 1), 1),0,0)
    if NLAB == len(labs):
        ax1.set_yticklabels(labs)
    #隐藏图像的底部边界
    ax1.spines['bottom'].set_visible(False)
    #plt.title('高密度排列图示')
    ###这个部分本身是不需要的，但是在视电阻率图加了colorbar之后，两个图大小一样了，因此这个排列图也加了一个colorbar
    divider = make_axes_locatable(ax1)
    cax = divider.append_axes('right', size="2%", pad=0.1)
    colormap = cm.get_cmap('jet')
    sm = plt.cm.ScalarMappable(cmap=colormap)
    cbar = plt.colorbar(sm, cax=cax)
    cbar.remove()           #不显示colorbar
    #####
    plt.draw()
    plt.pause(0.01)
    #plt.close()
    ##最终画出等值线图
    # plt.show()

    f = io.BytesIO()
    plt.savefig(f, format="png")
    print("___________________ERTPlot__________________22222")
    plt.close()
    return f.getvalue()

def ERTuiArrayABMN(num_electrodes0,Max_n_spacing0,Aspacing0,array_type0):
    # 开始采集之前，先计算排列，这里的参数不需要保存到数据文件里。
    #num_electrodes 电极总数（来自主机）
    #Max_n_spacing 最大隔离系数，若不设置，会根据电极总数计算，为 (num_electrodes - 3) （手部输入）
    #Aspacing 点距 （这个手部输入）
    #array_type 排列方式
    print("------------------ERTuiArray--------------------")
    num_electrodes = num_electrodes0
    Max_n_spacing = Max_n_spacing0
    Aspacing = Aspacing0
    array_type = array_type0
    print(num_electrodes)
    print(Max_n_spacing)
    print(Aspacing)
    print(array_type)
    # 最大隔离系数太大，系统自动调整
    if Max_n_spacing > num_electrodes - 3:
        print('最大隔离系数设置错误，隔离系数范围为' + str(1) + '到' + str(num_electrodes - 3))
        Max_n = num_electrodes - 3

    NS = np.sum(range(1, Max_n_spacing + 1)) + Max_n_spacing * (
            (num_electrodes - 3) - Max_n_spacing)  # 总排列数

    ABMN = np.zeros((NS, 4))  # ABMN电极排列
    c = 0
    # 计算排列电极位置 （这里假设只有一个ADC接收电极，当有多个ADC时，此部分再修改)
    if array_type == 0: #dipole-dipole
        for i in range(num_electrodes - 3):
            for j in range(i + 2, num_electrodes - 1):
                ABMN[c, 0] = i + 1  # A电极
                ABMN[c, 1] = i + 2  # B电极
                ABMN[c, 2] = j + 1  # M电极
                ABMN[c, 3] = j + 2  # N电极
                N = ABMN[c, 2] - ABMN[c, 1]
                if N > Max_n_spacing:
                    continue
                else:
                    c = c + 1
    elif array_type == 1:#pole-dipole
        for i in range(num_electrodes - 3):
            for j in range(i + 2, num_electrodes - 1):
                ABMN[c, 0] = 1  # A电极, 整个测量中不动
                ABMN[c, 1] = i + 2  # B电极
                ABMN[c, 2] = j + 1  # M电极
                ABMN[c, 3] = j + 2  # N电极
                N = ABMN[c, 2] - ABMN[c, 1]
                if N > Max_n_spacing:
                    continue
                else:
                    c = c + 1
    return str(ABMN)
